# koroFileHeader at Yungoal acer
# Create: 2023-01-12 11:44:34
# LastEdit: 2023-10-20 10:44:22
"""加密/解密
参考 ansible-vault 写的,安全有保证

被引用:
- env.py
"""
__author__ = '749B'

import base64
from binascii import hexlify
from binascii import unhexlify
import os
import sys

from Crypto.Cipher import AES
from Crypto.Hash import SHA256
from Crypto.Hash import HMAC
from Crypto.Util import Counter
from Crypto.Protocol.KDF import PBKDF2


HEADER = '$YUNGOAL_VAULT:'

def is_encrypted(data):
    """检查是否是加密的字符串"""
    return data.strip().startswith(HEADER)

def parse_vaulttext_envelope(vaulttext):
    """解析头部信息，返回密文部分"""
    return vaulttext.strip()[len(HEADER) + 1:]

def format_vaulttext_envelope(vaulttext):
    """添加头部信息"""
    return '\n'.join([HEADER, vaulttext])


def _gen_key_initctr(b_password, b_salt):
    """密码 + 盐 生成key"""
    key_length = 32
    iv_length = 16

    b_derivedkey = PBKDF2(
        b_password, b_salt, dkLen=(2 * key_length) + iv_length,
        count=10000, prf=lambda p, s: HMAC.new(p, s, SHA256).digest())
    b_iv = hexlify(b_derivedkey[(key_length * 2):(key_length * 2) + iv_length])

    b_key1 = b_derivedkey[:key_length]
    b_key2 = b_derivedkey[key_length:(key_length * 2)]

    return b_key1, b_key2, b_iv

def encrypt(plaintext, password):
    """加密"""
    if is_encrypted(plaintext):
        return plaintext

    b_salt = os.urandom(32)
    b_password = password.encode('utf-8')
    b_plaintext = plaintext.encode('utf-8')
    b_key1, b_key2, b_iv = _gen_key_initctr(b_password, b_salt)

    # _encrypt_pycrypto(b_plaintext, b_key1, b_key2, b_iv):
    bs = AES.block_size
    padding_length = (bs - len(b_plaintext) % bs) or bs
    b_plaintext += (padding_length * chr(padding_length)).encode('ascii')
    ctr = Counter.new(128, initial_value=int(b_iv, 16))
    cipher = AES.new(b_key1, AES.MODE_CTR, counter=ctr)
    b_ciphertext = cipher.encrypt(b_plaintext)
    hmac = HMAC.new(b_key2, b_ciphertext, SHA256)
    b_hmac =  hmac.hexdigest().encode('utf-8')

    b_vaulttext = b'\n'.join([hexlify(b_salt), b_hmac, hexlify(b_ciphertext)])
    b64_vaulttext = base64.encodebytes(b_vaulttext).decode('utf-8')
    fmt_vaulttext = format_vaulttext_envelope(b64_vaulttext)
    return fmt_vaulttext

def decrypt(fmt_vaulttext, password):
    """解密"""
    if not is_encrypted(fmt_vaulttext):
        return fmt_vaulttext

    b64_vaulttext = parse_vaulttext_envelope(fmt_vaulttext)
    b_password = password.encode('utf-8')

    # _parse_vaulttext(b_vaulttext):
    b_vaulttext = base64.decodebytes(b64_vaulttext.encode('utf-8'))
    b_salt, b_crypted_hmac, b_ciphertext = b_vaulttext.split(b'\n', 2)
    b_salt = unhexlify(b_salt)
    b_ciphertext = unhexlify(b_ciphertext)

    b_key1, b_key2, b_iv = _gen_key_initctr(b_password, b_salt)

    # _decrypt_pycrypto(cls, b_ciphertext, b_crypted_hmac, b_key1, b_key2, b_iv):
    hmac_decrypt = HMAC.new(b_key2, b_ciphertext, SHA256)
    if not _is_equal(b_crypted_hmac, hmac_decrypt.hexdigest().encode('utf-8')):
        return None

    ctr = Counter.new(128, initial_value=int(b_iv, 16))
    cipher = AES.new(b_key1, AES.MODE_CTR, counter=ctr)

    b_plaintext = cipher.decrypt(b_ciphertext)

    padding_length = b_plaintext[-1]
    b_plaintext = b_plaintext[:-padding_length]
    plaintext = b_plaintext.decode('utf-8')
    return plaintext

def _is_equal(b_a, b_b):
    """
    在恒定时间内比较两个 bytes 类型，防止时间攻击
    """
    if not (isinstance(b_a, bytes) and isinstance(b_b, bytes)):
        raise TypeError('_is_equal can only be used to compare two byte strings')

    # http://codahale.com/a-lesson-in-timing-attacks/
    if len(b_a) != len(b_b):
        return False

    result = 0
    for b_x, b_y in zip(b_a, b_b):
        result |= b_x ^ b_y
    return result == 0

def command_line_tool():
    """命令行工具：直接运行这个文件"""
    import getpass  # pylint: disable=import-outside-toplevel
    menu = {
        '1': "加密",
        '2': "解密",
        'q': "退出",
    }
    for k, v in menu.items():
        print(f'{k}\t{v}')
    choice = input("选择操作: ")
    if choice == 'q':
        print("退出")
        sys.exit(0)
    elif choice == '1':
        plaintext = input("请输入要加密的字符串: ")
        plaintext = plaintext.strip()
        password1 = getpass.getpass("加密密码: ")
        password2 = getpass.getpass("确认密码: ")
        if password1 != password2:
            print("确认密码错误")
            sys.exit(1)
        vaulttext = encrypt(plaintext, password1)
        print("加密字符串:")
        print(vaulttext)
    elif choice == '2':
        print("请输入要加密的字符串，新行开头 Ctrl+Z 回车结束: ")
        line = sys.stdin.readline()
        vaultlines = []
        while line:
            vaultlines.append(line)
            line = sys.stdin.readline()
        vaulttext = ''.join(vaultlines)
        password = getpass.getpass("加密密码: ")
        plaintext = decrypt(vaulttext, password)
        if not plaintext:
            print("解密失败。")
            sys.exit(1)
        print("解密字符串:")
        print(plaintext)
    else:
        print(f"不存在的选项({choice})，退出。")
        sys.exit(1)


if __name__ == '__main__':
    command_line_tool()
    